package pl.th.permission.plugin.node;

import static org.objectweb.asm.Opcodes.AASTORE;
import static org.objectweb.asm.Opcodes.ALOAD;
import static org.objectweb.asm.Opcodes.ANEWARRAY;
import static org.objectweb.asm.Opcodes.ARETURN;
import static org.objectweb.asm.Opcodes.DRETURN;
import static org.objectweb.asm.Opcodes.DUP;
import static org.objectweb.asm.Opcodes.FRETURN;
import static org.objectweb.asm.Opcodes.IFEQ;
import static org.objectweb.asm.Opcodes.IFNE;
import static org.objectweb.asm.Opcodes.INVOKESTATIC;
import static org.objectweb.asm.Opcodes.IRETURN;
import static org.objectweb.asm.Opcodes.LRETURN;
import static org.objectweb.asm.Opcodes.RETURN;

import org.objectweb.asm.AnnotationVisitor;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.Label;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.FrameNode;
import org.objectweb.asm.tree.InsnList;
import org.objectweb.asm.tree.InsnNode;
import org.objectweb.asm.tree.IntInsnNode;
import org.objectweb.asm.tree.JumpInsnNode;
import org.objectweb.asm.tree.LabelNode;
import org.objectweb.asm.tree.LdcInsnNode;
import org.objectweb.asm.tree.LineNumberNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.TypeInsnNode;
import org.objectweb.asm.tree.VarInsnNode;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.Set;

import pl.th.permission.plugin.Utils;
import pl.th.permission.plugin.node.annotation.AllAnnotation;
import pl.th.permission.plugin.node.annotation.AnnType;
import pl.th.permission.plugin.node.annotation.AnyAnnotation;
import pl.th.permission.plugin.node.annotation.CallbackAnnotation;
import pl.th.permission.plugin.node.annotation.NeedAnnotation;
import pl.th.permission.plugin.node.annotation.PartAnnotation;

public class PermissionMethodNode extends MethodNode {

    private NeedAnnotation needAnnotation;
    private AllAnnotation allGrantedAnnotation;
    private AllAnnotation allDeniedAnnotation;
    private AnyAnnotation anyGrantedAnnotation;
    private AnyAnnotation anyDeniedAnnotation;
    private PartAnnotation onGrantedAnnotation;
    private PartAnnotation onDeniedAnnotation;
    private CallbackAnnotation grantedCbAnnotation;
    private CallbackAnnotation deniedCbAnnotation;

    private int mFirstLineNumber = -1;
    private LabelNode mReturnLabelNode;

    public PermissionMethodNode(final int access, final String name, final String descriptor,
                                final String signature, final String[] exceptions) {
        super(Opcodes.ASM7, access, name, descriptor, signature, exceptions);
    }

    public NeedAnnotation getNeedAnnotation() {
        return needAnnotation;
    }

    public AllAnnotation getAllGrantedAnnotation() {
        return allGrantedAnnotation;
    }

    public AllAnnotation getAllDeniedAnnotation() {
        return allDeniedAnnotation;
    }

    public AnyAnnotation getAnyGrantedAnnotation() {
        return anyGrantedAnnotation;
    }

    public AnyAnnotation getAnyDeniedAnnotation() {
        return anyDeniedAnnotation;
    }

    public PartAnnotation getOnGrantedAnnotation() {
        return onGrantedAnnotation;
    }

    public PartAnnotation getOnDeniedAnnotation() {
        return onDeniedAnnotation;
    }

    public CallbackAnnotation getGrantedCbAnnotation() {
        return grantedCbAnnotation;
    }

    public CallbackAnnotation getDeniedCbAnnotation() {
        return deniedCbAnnotation;
    }

    @Override
    public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
        if (visible) {
            if (visibleAnnotations == null) {
                visibleAnnotations = new ArrayList<>(1);
            }
            CallbackAnnotation node = null;
            if (Utils.isAnnotateMethodByNA(descriptor)) {
                needAnnotation = new NeedAnnotation();
                node = needAnnotation;
            } else if (Utils.isAnnotatedByAllGranted(descriptor)) {
                allGrantedAnnotation = new AllAnnotation(AnnType.ON_ALL_PERMISSION_GRANTED);
                node = allGrantedAnnotation;
            } else if (Utils.isAnnotatedByAllDenied(descriptor)) {
                allDeniedAnnotation = new AllAnnotation(AnnType.ON_ALL_PERMISSION_DENIED);
                node = allDeniedAnnotation;
            } else if (Utils.isAnnotatedByAnyDenied(descriptor)) {
                anyDeniedAnnotation = new AnyAnnotation(AnnType.ON_ANY_PERMISSION_DENIED);
                node = anyDeniedAnnotation;
            } else if (Utils.isAnnotatedByAnyGranted(descriptor)) {
                anyGrantedAnnotation = new AnyAnnotation(AnnType.ON_ANY_PERMISSION_GRANTED);
                node = anyGrantedAnnotation;
            } else if (Utils.isAnnotatedByGrantedCallback(descriptor)) {
                grantedCbAnnotation = new CallbackAnnotation(AnnType.ON_GRANTED_CALLBACK);
                node = grantedCbAnnotation;
            } else if (Utils.isAnnotatedByDeniedCallback(descriptor)) {
                deniedCbAnnotation = new CallbackAnnotation(AnnType.ON_DENIED_CALLBACK);
                node = deniedCbAnnotation;
            } else if (Utils.isAnnotateMethodByG(descriptor)) {
                onGrantedAnnotation = new PartAnnotation(AnnType.ON_PERMISSION_GRANTED);
                node = onGrantedAnnotation;
            } else if (Utils.isAnnotateMethodByD(descriptor)) {
                onDeniedAnnotation = new PartAnnotation(AnnType.ON_PERMISSION_DENIED);
                node = onDeniedAnnotation;
            }
            if (node != null) {
                node.checkUsedRight(desc);
                visibleAnnotations.add(node);
                return node;
            }
        }
        return super.visitAnnotation(descriptor, visible);
    }

    @Override
    public void visitInsn(int opcode) {
        switch (opcode) {
            case RETURN:
            case IRETURN:
            case FRETURN:
            case ARETURN:
            case LRETURN:
            case DRETURN:
                if (shouldInjectCodeAtEnd()) {
                    addPermissionInsnList();
                }
                break;
        }
        super.visitInsn(opcode);
    }

    @Override
    public void visitMaxs(int maxStack, int maxLocals) {
        if (needAnnotation != null) {
            if (shouldInjectCodeAtBegin()) {
                maxStack = Math.max(maxStack, 5);
            } else {
                maxStack += 5;
            }
            maxLocals = Math.max(1, maxLocals);
        }
        super.visitMaxs(maxStack, maxLocals);
    }

    @Override
    public void accept(ClassVisitor classVisitor) {
        if (shouldInjectCodeAtBegin()) {
            findNeededLabelNode();
            insertPermissionInsnList();
        }
        super.accept(classVisitor);
    }

    private boolean shouldInjectCodeAtBegin() {
        return needAnnotation != null && allGrantedAnnotation != null
                && needAnnotation.getRequestCode() == allGrantedAnnotation.getRequestCode();
    }

    private boolean shouldInjectCodeAtEnd() {
        return needAnnotation != null && (allGrantedAnnotation == null
                || needAnnotation.getRequestCode() != allGrantedAnnotation.getRequestCode());
    }

    private void findNeededLabelNode() {
        final Iterator<AbstractInsnNode> insnList = instructions.iterator();
        while (insnList.hasNext()) {
            AbstractInsnNode insnNode = insnList.next();
            if (insnNode instanceof LabelNode) {
                mReturnLabelNode = (LabelNode) insnNode;
            }
            if (insnNode instanceof LineNumberNode && mFirstLineNumber == -1) {
                mFirstLineNumber = ((LineNumberNode)insnNode).line;
            }
            if (insnNode instanceof InsnNode && insnNode.getOpcode() == RETURN) {
                break;
            }
        }
    }

    // inject code of request permission at the start of method
    private void insertPermissionInsnList() {
        InsnList insertedList = new InsnList();
        final NeedAnnotation annotation = needAnnotation;
        if (mFirstLineNumber > 0) {
            LabelNode l0 = getLabelNode(new Label());
            insertedList.add(l0);
            insertedList.add(new LineNumberNode(mFirstLineNumber, l0));
        }
        push(insertedList, annotation.getRequestCode());
        Set<String> permissions = annotation.getPermissions();
        push(insertedList, permissions.size());
        insertedList.add(new TypeInsnNode(ANEWARRAY, "java/lang/String"));
        int i = 0;
        for (String permission : permissions) {
            insertedList.add(new InsnNode(DUP));
            push(insertedList, i);
            insertedList.add(new LdcInsnNode(permission));
            insertedList.add(new InsnNode(AASTORE));
            i++;
        }
        insertedList.add(new LdcInsnNode(annotation.getClzType()));
        insertedList.add(new VarInsnNode(ALOAD, 0));
        insertedList.add(new MethodInsnNode(INVOKESTATIC, Utils.NEED_INJECT_A_METHOD_OWNER,
                Utils.NEED_INJECT_A_METHOD_NAME, Utils.NEED_INJECT_A_METHOD_DESC, false));
        if (mReturnLabelNode != null) {
            insertedList.add(new JumpInsnNode(IFEQ, mReturnLabelNode));
        } else {
            /* if mReturnLabelNode == null and try...catch... is used in the annotated method body,
               it will build failed */
            LabelNode l1 = getLabelNode(new Label());
            insertedList.add(new JumpInsnNode(IFNE, l1));
            LabelNode l2 = getLabelNode(new Label());
            insertedList.add(l2);
            insertedList.add(new InsnNode(RETURN));
            insertedList.add(l1);
            insertedList.add(new FrameNode(Opcodes.F_SAME, 0, null, 0, null));
        }
        instructions.insert(insertedList);
    }

    // inject code of request permission at the end of method
    private void addPermissionInsnList() {
        final NeedAnnotation annotation = needAnnotation;
        Label label0 = new Label();
        super.visitLabel(label0);
        push(annotation.getRequestCode());
        Set<String> permissions = annotation.getPermissions();
        push(permissions.size());
        super.visitTypeInsn(ANEWARRAY, "java/lang/String");
        int i = 0;
        for (String permission : permissions) {
            super.visitInsn(DUP);
            push(i);
            super.visitLdcInsn(permission);
            super.visitInsn(AASTORE);
            i++;
        }
        super.visitLdcInsn(annotation.getClzType());
        super.visitVarInsn(ALOAD, 0);
        super.visitMethodInsn(INVOKESTATIC, Utils.NEED_INJECT_A_METHOD_OWNER,
                Utils.NEED_INJECT_A_CALLBACK_METHOD_NAME, Utils.NEED_INJECT_A_CALLBACK_METHOD_DESC, false);
    }

    private void push(final int value) {
        if (value >= -1 && value <= 5) {
            super.visitInsn(Opcodes.ICONST_0 + value);
        } else if (value >= Byte.MIN_VALUE && value <= Byte.MAX_VALUE) {
            super.visitIntInsn(Opcodes.BIPUSH, value);
        } else if (value >= Short.MIN_VALUE && value <= Short.MAX_VALUE) {
            super.visitIntInsn(Opcodes.SIPUSH, value);
        } else {
            super.visitLdcInsn(value);
        }
    }

    private void push(final InsnList insnList, final int value) {
        if (value >= -1 && value <= 5) {
            insnList.add(new InsnNode(Opcodes.ICONST_0 + value));
        } else if (value >= Byte.MIN_VALUE && value <= Byte.MAX_VALUE) {
            insnList.add(new IntInsnNode(Opcodes.BIPUSH, value));
        } else if (value >= Short.MIN_VALUE && value <= Short.MAX_VALUE) {
            insnList.add(new IntInsnNode(Opcodes.SIPUSH, value));
        } else {
            insnList.add(new LdcInsnNode(value));
        }
    }
}
